package com.wry.wry_test.service.impl;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @Author: wangrongyi
 * @Date: 2024/7/9 14:33
 * @Description:
 */
@RestController
@RequestMapping("/chatGPT")
@Slf4j
public class ChatGPTServiceController {


    @Value("${openai.api.key:sk-proj-WyrAAUWvTUgjp8M5YaRcT3BlbkFJf6OfwgaRnhyX1dGoBmbR}")
    private String apiKey;

    @Resource
    private RestTemplate restTemplate;


    @GetMapping("/chat")
    public String chat(@RequestParam("prompt") String prompt) {
        return this.getChatGPTResponse(prompt);
    }


    private String getChatGPTResponse(String prompt) {
        log.info("ChatGPT request: {}", prompt);
        String url = "https://api.openai.com/v1/engines/davinci-codex/completions";

        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + "sk-proj-WyrAAUWvTUgjp8M5YaRcT3BlbkFJf6OfwgaRnhyX1dGoBmbR");
        headers.set("Content-Type", "application/json");

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("prompt", prompt);
        requestBody.put("max_tokens", 150);

        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

        ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, entity, Map.class);

        Map<String, Object> responseBody = response.getBody();
        if (responseBody != null && responseBody.containsKey("choices")) {
            Map<String, Object> firstChoice = ((List<Map<String, Object>>) responseBody.get("choices")).get(0);
            return (String) firstChoice.get("text");
        }

        return null;
    }

}